﻿#nullable enable

namespace NServiceBus;

using System;
using System.Threading.Tasks;
using Extensibility;
using Pipeline;

class PopulateAutoCorrelationHeadersForRepliesBehavior : IBehavior<IOutgoingReplyContext, IOutgoingReplyContext>
{
    public Task Invoke(IOutgoingReplyContext context, Func<IOutgoingReplyContext, Task> next)
    {
        FlowDetailsForRequestingSagaToOutgoingMessage(context);

        return next(context);
    }

    static void FlowDetailsForRequestingSagaToOutgoingMessage(IOutgoingReplyContext context)
    {
        if (context.TryGetIncomingPhysicalMessage(out var incomingMessage))
        {
            incomingMessage.Headers.TryGetValue(Headers.OriginatingSagaId, out var sagaId);
            incomingMessage.Headers.TryGetValue(Headers.OriginatingSagaType, out var sagaType);

            if (context.GetOperationProperties().TryGet<State>(out var state))
            {
                sagaId = state.SagaIdToUse;
                sagaType = state.SagaTypeToUse;
            }

            if (!string.IsNullOrEmpty(sagaId))
            {
                context.Headers[Headers.SagaId] = sagaId;
            }

            if (!string.IsNullOrEmpty(sagaType))
            {
                context.Headers[Headers.SagaType] = sagaType;
            }
        }
    }

    public class State
    {
        public string? SagaIdToUse { get; init; }

        public string? SagaTypeToUse { get; init; }
    }
}